import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms


class MedicalImageDataset(Dataset):
    def __init__(self, list_file):
        """
        初始化数据集
        :param list_file: 包含图像和标签路径的列表文件
        """
        self.data = []
        self.root_dir = '/'.join(list_file.split('/')[:-1])
        i = 0
        with open(list_file, 'r') as file:
            for line in file:
                image_path, label_path = line.strip().split()  # 假设每行两个路径由空格分隔
                image_path = self.root_dir + '/' + image_path
                label_path = self.root_dir + '/' + label_path
                self.data.append((image_path, label_path))
                # i += 1
                # if i % 50 == 0:
                #     break

    def __len__(self):
        """
        返回数据集中的样本总数
        """
        return len(self.data)

    def __getitem__(self, idx):
        """
        根据索引 idx 返回一个样本
        :param idx: 样本索引
        :return: image 和 mask 的元组
        """
        img_path, label_path = self.data[idx]

        # 加载 numpy 数组
        image = np.load(img_path)
        mask = np.load(label_path)

        # 将 numpy 数组转换为 torch tensors
        image = torch.from_numpy(image).float()
        mask = torch.from_numpy(mask).long()
        # image = transforms.Resize((256, 256))(transforms.ToTensor()(image))
        # mask = transforms.Resize((256, 256))(transforms.ToTensor()(mask))

        return image, mask


if __name__ == '__main__':
    # 指定列表文件的位置
    list_file = '/Volumes/For_Mac/dateset/Synapse_npy/train_list.txt'

    # 创建数据集实例
    dataset = MedicalImageDataset(list_file)

    # 创建 DataLoader
    dataloader = DataLoader(dataset, batch_size=1, shuffle=True, num_workers=2)
    # 示例：迭代一次数据加载器，打印第一批图像和标签的形状
    for images, masks in dataloader:

        print(f"Images batch shape: {images.shape}")
        print(f"Masks batch shape: {masks.shape}")
